feat(jax): freeze models with hessian output#5608
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a ChangesJAX freeze Hessian support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant freeze
participant serialization
participant model
participant HLO
CLI->>freeze: freeze --hessian
freeze->>serialization: deserialize_to_file(hessian=True)
serialization->>model: enable_hessian()
serialization->>serialization: set model_def_script hessian_mode
serialization-->>HLO: export Hessian-enabled model
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/jax/test_training.py (1)
199-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert
hessianin the dispatcher call.This still passes if
main()dropsargs.hessianbefore invokingfreeze(). Please assert the patched call carries the new flag as well.Suggested test tightening
main(args) freeze_entrypoint.assert_called_once() + self.assertIn("hessian", freeze_entrypoint.call_args.kwargs) + self.assertFalse(freeze_entrypoint.call_args.kwargs["hessian"])🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/jax/test_training.py` around lines 199 - 214, The test for the JAX CLI freeze dispatch is too loose because it only checks that the patched freeze entrypoint was called. Tighten test_main_dispatches_freeze in test_training.py by asserting the call includes the hessian flag from the argparse.Namespace, so main() must forward args.hessian into deepmd.jax.entrypoints.main.freeze rather than dropping it. Use the existing freeze_entrypoint mock and verify its call arguments reflect the new flag.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/jax/test_training.py`:
- Around line 199-214: The test for the JAX CLI freeze dispatch is too loose
because it only checks that the patched freeze entrypoint was called. Tighten
test_main_dispatches_freeze in test_training.py by asserting the call includes
the hessian flag from the argparse.Namespace, so main() must forward
args.hessian into deepmd.jax.entrypoints.main.freeze rather than dropping it.
Use the existing freeze_entrypoint mock and verify its call arguments reflect
the new flag.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 44d1886d-fd38-4bbe-a056-c56edc3c92fc
📒 Files selected for processing (7)
deepmd/jax/entrypoints/freeze.pydeepmd/jax/infer/deep_eval.pydeepmd/jax/jax2tf/serialization.pydeepmd/jax/model/hlo.pydeepmd/jax/utils/serialization.pydeepmd/main.pysource/tests/jax/test_training.py
njzjz
force-pushed
the
feat/jax-hessian-freeze
branch
from
16058df to
6eabd63
Compare
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5608 +/- ##
==========================================
- Coverage 79.65% 79.52% -0.14%
==========================================
Files 1015 1015
Lines 115740 115805 +65
Branches 4275 4272 -3
==========================================
- Hits 92198 92096 -102
- Misses 22001 22167 +166
- Partials 1541 1542 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
I'd like to wait for #5613. |
njzjz
force-pushed
the
feat/jax-hessian-freeze
branch
from
a476689 to
30735c1
Compare
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
| OutputVariableCategory.REDU, | ||
| OutputVariableCategory.DERV_R, | ||
| OutputVariableCategory.DERV_C_REDU, | ||
| OutputVariableCategory.DERV_R_DERV_R, |
There was a problem hiding this comment.
Adding DERV_R_DERV_R here without a get_has_hessian() guard makes every JAX model request energy_derv_r_derv_r, including ones frozen without hessian. self.output_def is the consumer def from DeepPot.output_def, which hard-codes r_hessian=True, so this category is always present (it is not gated on hessian_mode). For a non-hessian model the stablehlo/savedmodel never produces that key, so _eval_model falls into the fill branch and allocates np.full([nframes, 3*natoms, 3*natoms], nan) — an O(N^2) array (~72 MB at 1000 atoms) on the common non-atomic force/energy path — which DeepPot.eval then discards because get_has_hessian() is False.
This reintroduces the regression fixed for the PT backend in #5045 ("fix: remove hessian outdef if not necessary", 87cb6ef). The PR already adds get_has_hessian() (below), so the fix is to mirror pt (deepmd/pt/infer/deep_eval.py) and drop DERV_R_DERV_R from the request defs when not self.get_has_hessian().
There was a problem hiding this comment.
Fixed in bab82b1 by filtering DERV_R_DERV_R from JAX request definitions whenever get_has_hessian() is false, matching the PyTorch backend behavior. Added regression coverage for standard non-Hessian models.
Validation:
pytest source/tests/jax/test_training.py -k "deep_eval_requests_hessian or deep_eval_skips_hessian or main_dispatches_freeze or hlo_hessian_mode" -vruff check .ruff format --check .- commit pre-hooks (including pylint) passed
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Avoid requesting Hessian output from standard frozen JAX models and add regression coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Summary
dp freeze --hessiansupport for the JAX backend.jax,.hlo, and.savedmodelserializationTests
source venv/bin/activate && pytest source/tests/jax/test_training.py::TestJAXTraining::test_freeze_entrypoint_uses_checkpoint_pointer source/tests/jax/test_training.py::TestJAXTraining::test_main_dispatches_freeze source/tests/jax/test_training.py::TestJAXTraining::test_hlo_hessian_mode_updates_output_def source/tests/jax/test_training.py::TestJAXTraining::test_deep_eval_requests_hessian_for_hessian_model -qsource venv/bin/activate && ruff check .source venv/bin/activate && ruff format .Summary by CodeRabbit
New Features
--hessianoption to the model freezing command to include Hessian information in exported outputs.Bug Fixes
Tests